#include <iostream>
#include <iomanip>
#include <vector>
#include <cstdlib>
#include <complex>
#include <omp.h>
#include <fft/fft.h>
#include <mambo.hpp>
using namespace std;

/* driver to test fft libray
 *
 * Need compiler macro: -DMYFFT=0 
 * * -DMYFFT=0 with fftw
 * * -DMYFFT=1 with essl
 * * -DMYFFT=2 with mkl
 */
int main(int argc, char** argv)
{

  ///1D FFT dimension
  int fft_dim=12288;
  ///x grid
  int npx=512;
  ///y grid
  int npy=512;
  ///number of ffts to be performed per thread
  int num_fft_do=2;
  ///number of iterations
  int niters=1;
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size()) //get fft dimension
      fft_dim=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //get how many fft to perform
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //get how many fft to perform
      npy=atoi(argv[++ic]);
    else if(a.find("-d")<a.size()) //number of fftws
      num_fft_do=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-h")<a.size())
    {
      cerr << "Usage: -s fft-size -npx x-grid -npy y-grid -d fft-to-do -i iterations" << endl;
    }
    ++ic;
  }

#if defined(FFT_R2C)
  int num_fft=fft_dim/npx*fft_dim/npy;
#else
  int num_fft=fft_dim/npx*fft_dim/npy/2;
#endif
  num_fft_do=std::min(num_fft_do,num_fft);

  double flops_c=(log(static_cast<double>(fft_dim))/log(2.))*5.*fft_dim;
  double flops=flops_c;

  //typedefs for the real and complex
  typedef double real_type;
  typedef complex<double> complex_type;

  ///FFT engine which can do multiple ffts of domany==true
  const bool domany=true;
  int np=omp_get_max_threads();

  cout << "BW FFT dim= " << fft_dim 
    << " nx_loc " << fft_dim/npx
    << " ny_loc " << fft_dim/npy
    << " nz_loc " << fft_dim
    << "\nBW target ffts = " <<num_fft
    << " fft done = " << num_fft_do
    << " iterations= " << niters ;

#if defined(FFT_R2C)
  //real-to-complex transforms
  flops *=0.5;
  typedef FFT_DFT<real_type,real_type,domany,MYFFT> fft_eng_type;
  int fft_buffer=fft_dim+2;
  real_type* target = new real_type[np*fft_buffer*num_fft];
  cout << " FLOPS r2c= " << flops<< endl;
#else
  //complex-to-complex transforms
  typedef FFT_DFT<real_type,complex_type,domany,MYFFT> fft_eng_type;
  int fft_buffer=fft_dim;
  complex_type* target = new complex_type[np*fft_buffer*num_fft];
  cout << " FLOPS c2c= " << flops<< endl;
#endif

  //create engines per thread but ignore thread-local allocators
  vector<fft_eng_type*> fft_doer(np);
  for(int ip=0; ip<np; ++ip)
  {
    int offset=ip*fft_buffer*num_fft;
    fft_doer[ip]=new fft_eng_type(fft_dim,target+offset,num_fft_do);
  }

  cout << "PHASEC" << endl;
  mambo::begin();
  Timer wtimer;
//#pragma omp parallel 
  {
    //int ip=omp_get_thread_num();
    int ip=0;
    //local address [first,last)
    int first=ip*fft_buffer*num_fft;
    int last=(ip+1)*fft_buffer*num_fft;
    fft_eng_type& myfft(*fft_doer[ip]);
    for(int iter=0,offset=first; iter<niters; ++iter)
    {
      myfft.transform_forward(target+offset,num_fft_do);
      myfft.transform_backward(target+offset,num_fft_do);
      offset+=fft_buffer*num_fft_do;
      if(offset>=last) offset=first;//rotate buffer
    }
  }
  double dt=wtimer.elapsed();

  flops *= num_fft_do*2*niters*1e-9;
  cout << "BW avgGFLOPS " << flops/dt 
    <<" totGFLOPS " << flops << " wtime " << dt << endl;

  mambo::end();
  cout << "PHASED" << endl;
  return 0;
}

